package top.yudoge.dogeeditor.llm.embedding

import cn.hutool.core.io.FileUtil
import dev.langchain4j.data.embedding.Embedding
import dev.langchain4j.data.segment.TextSegment
import dev.langchain4j.store.embedding.EmbeddingSearchRequest
import dev.langchain4j.store.embedding.EmbeddingSearchResult
import dev.langchain4j.store.embedding.EmbeddingStore
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore
import org.springframework.beans.factory.annotation.Value
import org.springframework.stereotype.Component
import java.io.File
import java.nio.charset.StandardCharsets

@Component
class EmbeddingStoreHolder {

    @Value("\${de.app.embedding.delta-id-splitor:[SPL]}")
    private lateinit var deltaIdSplitor: String

    @Value("\${de.app.embedding.chroma.endpoint:http://localhost:8000}")
    private lateinit var endpoint: String

    @Value("\${de.app.embedding.delta-file-path}")
    private lateinit var deltaFilePath: String

    private val storeMap = mutableMapOf<String, DeltaEmbeddingStore>()

    private val collectionKeyGenerator = EmbeddingCollectionKeyGenerator()


    fun getStore(collection: String): DeltaEmbeddingStore {
        synchronized(this) {
            if (!storeMap.containsKey(collection)) {
                var collectionKey = collectionKeyGenerator.key(collection)
                val embeddingStore: EmbeddingStore<TextSegment>? = ChromaEmbeddingStore
                    .builder()
                    .baseUrl(endpoint)
                    .collectionName(collectionKey)
                    .logRequests(true)
                    .logResponses(true)
                    .build()

                storeMap[collection] = DeltaEmbeddingStore(deltaFilePath, embeddingStore!!)
            }
            return storeMap[collection]!!
        }
    }

    fun getStore(dsIdentifier: String, dbName: String, schemaName: String): DeltaEmbeddingStore {
        val collection = "$dsIdentifier$deltaIdSplitor$dbName$deltaIdSplitor$schemaName"
        return getStore(collection)
    }


    fun genId(dsIdentifier: String, dbName: String, schemaName: String, tableName: String): String {
        return "$dsIdentifier$deltaIdSplitor$dbName$deltaIdSplitor$schemaName${deltaIdSplitor}$tableName"
    }

    fun getTableName(id: String): String {
        return id.split(deltaIdSplitor).last()
    }

    open class DeltaEmbeddingStore(
        private val deltaFilePath: String,
        private val store: EmbeddingStore<TextSegment>
    ) {

        /**
         * 向向量库中添加embedding，内容为segment，id为id
         */
        fun add(embedding: Embedding, segment: TextSegment, id: String) {
            store.remove(id)
            store.addAll(
                listOf(id),
                listOf(embedding),
                listOf(segment)
            )

            val file = File(deltaFilePath, id)

            FileUtil.writeString(
                segment.text(),
                file,
                StandardCharsets.UTF_8,
            )
        }

        /**
         * 根据id判断向量库中当前的segment和传入的segment是否一致，不一致则返回true，一致则返回true，向量库中不存在指定id，返回true，否则返回false
         */
        fun hasModifications(id: String, segment: TextSegment): Boolean {
            val file = File(deltaFilePath, id)
            if (!file.exists()) {
                return true
            }

            var oldText = FileUtil.readString(
                file,
                StandardCharsets.UTF_8,
            )
            return oldText != segment.text()
        }

        fun search(embeddingSearchRequest: EmbeddingSearchRequest): EmbeddingSearchResult<TextSegment> {
            return store.search(embeddingSearchRequest)
        }

    }

}